import random
from numpy.ma import *


def load_dataset():
    posting_list = [['my', 'dog', 'has', 'flea', 'problems', 'help', 'please'],
                   ['maybe', 'not', 'take', 'him', 'to', 'dog', 'park', 'stupid'],
                   ['my', 'dalmation', 'is', 'so', 'cute', 'I', 'love', 'him'],
                   ['stop', 'posting', 'stupid', 'worthless', 'garbage'],
                   ['mr', 'licks', 'ate', 'my', 'steak', 'how', 'to', 'stop', 'him'],
                   ['quit', 'buying', 'worthless', 'dog', 'food', 'stupid']]
    class_vec = [0, 1, 0, 1, 0, 1]  # 1 is abusive, 0 not
    return posting_list, class_vec


def create_vocab_list(dataset):
    vocabset = set([])
    for document in dataset:
        #取并集
        vocabset = vocabset | set(document)
    return list(vocabset)


def set_of_words2_vec(vocab_list, input_set):
    return_vec = [0] * len(vocab_list)
    for word in input_set:
        if word in vocab_list:
            return_vec[vocab_list.index(word)] = 1
    return return_vec


def train_nb0(train_matri, train_cate):
    num_traindoc = len(train_matri)
    num_words = len(train_matri[0])
    p_abusive = sum(train_cate) / float(num_traindoc)

    # p0_num = zeros(num_words)
    # p1_num = zeros(num_words)
    #拉普拉斯平滑，λ取1
    p0_num = ones(num_words)
    p1_num = ones(num_words)

    #λ乘以类别数
    p0_denom = 2.0
    p1_denom = 2.0
    for i in range(num_traindoc):
        if train_cate[i] == 1:
            p1_num += train_matri[i]
            p1_denom += sum(train_matri[i])
        else:
            p0_num += train_matri[i]
            p0_denom += sum(train_matri[i])
    p1_vec = log(p1_num / p1_denom)
    p0_vec = log(p0_num / p0_denom)
    return p0_vec, p1_vec, p_abusive


def classify_nb(vec2_classify, p0_vec, p1_vec, p_class1):
    p1 = sum(vec2_classify * p1_vec) + log(p_class1)
    p0 = sum(vec2_classify * p0_vec) + log(1-p_class1)
    if p1 > p0:
        return 1
    else:
        return 0


def testing_nb():
    test_list, class_vec = load_dataset()
    vocab_list = create_vocab_list(test_list)
    train_matri = []
    for doc in test_list:
        train_matri.append(set_of_words2_vec(vocab_list, doc))
    print(train_matri)
    p0_vec, p1_vec, p_abusive = train_nb0(train_matri, class_vec)
    test_entry = ['love', 'my', 'dalmation']
    test_entry1 = ['stupid', 'garbage', 'kill']
    class_label = classify_nb(array(set_of_words2_vec(vocab_list, test_entry1)), p0_vec,
                p1_vec, p_abusive)
    if class_label:
        print("Abusive")
    else:
        print("Not")


def bag_of_words2_vec(vocab_list, input_set):
    ret_matri = [0] * len(vocab_list)
    for word in input_set:
        if word in vocab_list:
            ret_matri[vocab_list.index(word)] += 1
    return ret_matri


def text_parse(big_string):
    import re
    list_of_tokens = re.split(r'\W*', big_string)
    return [token.lower() for token in list_of_tokens if len(token) > 2]


def spam_test():
    doc_list = []
    class_list = []
    full_text = []
    for i in range(1, 26):
        word_list = text_parse(open('email/spam/%d.txt' % i).read())
        doc_list.append(word_list)
        full_text.extend(word_list)
        class_list.append(1)
        word_list = text_parse(open('email/ham/%d.txt' % i).read())
        doc_list.append(word_list)
        full_text.extend(word_list)
        class_list.append(0)
    vocab_list = create_vocab_list(doc_list)
    train_set = list(range(50))
    test_set = []
    for i in range(10):
        rand_index = int(random.uniform(0, len(train_set)))
        test_set.append(rand_index)
        del(train_set[rand_index])
    train_matri = []
    train_class = []
    for doc_index in train_set:
        train_matri.append(set_of_words2_vec(vocab_list, doc_list[doc_index]))
        train_class.append(class_list[doc_index])
    p0_vec, p1_vec, p_spam = train_nb0(array(train_matri), array(train_class))
    error_count = 0
    for doc_index in test_set:
        word_vec = set_of_words2_vec(vocab_list, doc_list[doc_index])
        if classify_nb(array(word_vec), p0_vec, p1_vec, p_spam) !=class_list[doc_index]:
            error_count += 1
    print('The error rate is: ', float(error_count)/len(test_set))


if __name__ == '__main__':

    for i in range(10):
        spam_test()
